package com.xr.system.common.aspect;

import cn.hutool.core.util.StrUtil;
import com.xr.common.utils.IpUtil;
import com.xr.system.common.annotation.TrafficLimit;
import com.xr.system.common.enums.LimitTypeEnum;
import com.xr.system.common.exception.ServiceException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Collections;

@Aspect
@Component
@Slf4j
public class TrafficLimitAspect {

	@Autowired
	private RedisTemplate<String, Object> redisTemplate;

	@Pointcut("@annotation(trafficLimit)")
	public void limitPointCut(TrafficLimit trafficLimit) {
	}

	@Around(value = "limitPointCut(trafficLimit)", argNames = "pjp,trafficLimit")
	public Object interceptor(ProceedingJoinPoint pjp, TrafficLimit trafficLimit) {
		MethodSignature signature = (MethodSignature) pjp.getSignature();
		Method method = signature.getMethod();
		LimitTypeEnum limitType = trafficLimit.limitType();
		String key;
		int limitPeriod = trafficLimit.period();
		int limitCount = trafficLimit.count();
		// 根据限流类型获取不同的key ,如果不传我们会以方法名作为key
		switch (limitType) {
		case IP:
			key = IpUtil.getIp();
			break;
		case CUSTOM:
			key = trafficLimit.key();
			break;
		default:
			key = method.getName();
		}
		try {
			String luaScript = buildLuaScript();
			RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
			Number count = redisTemplate.execute(redisScript,
					Collections.singletonList(StrUtil.join(trafficLimit.prefix(), key)), limitCount, limitPeriod);
			log.info("Access try count is {}  and key = {}", count, key);
			if (count != null && count.intValue() <= limitCount) {
				return pjp.proceed();
			}
			else {
				throw new ServiceException("服务器繁忙，请稍后再试！");
			}
		}
		catch (Throwable e) {
			if (e instanceof RuntimeException) {
				throw new RuntimeException(e.getLocalizedMessage());
			}
			throw new RuntimeException("server exception");
		}
	}

	/**
	 * 编写 redis Lua 限流脚本
	 */
	public String buildLuaScript() {
		return "local c" + "\nc = redis.call('get',KEYS[1])" +
		// 调用不超过最大值，则直接返回
				"\nif c and tonumber(c) > tonumber(ARGV[1]) then" + "\nreturn c;" + "\nend" +
				// 执行计算器自加
				"\nc = redis.call('incr',KEYS[1])" +
				// 从第一次调用开始限流，设置对应键值的过期
				"\nif tonumber(c) == 1 then" + "\nredis.call('expire',KEYS[1],ARGV[2])" + "\nend" + "\nreturn c;";
	}

}
